#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
增强版命名实体识别处理器
结合词典匹配、正则表达式和预训练模型进行实体识别
"""

import re
import logging
import os
from typing import List, Dict, Any

# 配置日志
# 确保日志目录存在
log_dir = "logs"
if not os.path.exists(log_dir):
    os.makedirs(log_dir)

# 为这个模块创建独立的logger
logger = logging.getLogger('advanced_ner')
logger.setLevel(logging.INFO)

# 清除可能存在的旧处理器
for handler in logger.handlers[:]:
    logger.removeHandler(handler)

# 创建文件处理器
file_handler = logging.FileHandler(os.path.join(log_dir, "advanced_ner.log"), encoding='utf-8')
file_handler.setLevel(logging.INFO)

# 创建控制台处理器
console_handler = logging.StreamHandler()
console_handler.setLevel(logging.INFO)

# 创建格式化器
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
file_handler.setFormatter(formatter)
console_handler.setFormatter(formatter)

# 添加处理器到logger
logger.addHandler(file_handler)
logger.addHandler(console_handler)

# 防止日志传播到根logger
logger.propagate = False

class AdvancedNERProcessor:
    """
    增强版命名实体识别处理器
    结合词典匹配和正则表达式进行实体识别
    """
    
    def __init__(self):
        """
        初始化处理器，设置词典和正则表达式模式
        """
        # 预定义词典
        self.dictionaries = {
            "组织": ["承德钒钛", "河钢集团", "首钢", "宝钢集团"],
            "工种": ["钳工", "焊工", "技术工程师", "电工"],
            "证书": ["操作证", "焊工证", "电工证"],
            "单位": ["MPa", "℃"],
            "方案": ["行动工作方案"],
            "职务": ["工程师", "经理", "主任"],
            "学历": ["本科", "硕士", "博士"],
            "姓名": ["张三", "李四", "王五"],
            "岗位": ["技术岗", "管理岗", "操作岗"]
        }
        
        # 领域特定词典
        self.domain_dictionaries = {
            "manufacturing": {
                "组织": ["承德钒钛", "河钢集团", "首钢", "宝钢集团", "鞍钢", "武钢"],
                "工种": ["钳工", "焊工", "技术工程师", "电工", "车工", "铣工"],
                "证书": ["操作证", "焊工证", "电工证", "资格证书"],
                "压力值": ["MPa", "kPa", "Pa"],
                "温度值": ["℃"]
            },
            "enterprise": {
                "方案": ["行动工作方案"],
                "部门": ["人力资源部", "财务部", "技术部", "生产部"],
                "百分比": ["%"],
                "日期": ["年", "月", "日"]
            }
        }
        
        # 正则表达式模式
        self.patterns = {
            "压力值": r"(\d+\.?\d*)\s*(MPa|kPa|Pa)",
            "温度值": r"(\d+\.?\d*)\s*(℃|度)",
            "百分比": r"(\d+\.?\d*)\s*%",
            "日期": r"(\d{4}年\d{1,2}月\d{1,2}日|\d{4}-\d{1,2}-\d{1,2}|\d{4}/\d{1,2}/\d{1,2})",
            "证书编号": r"[A-Z0-9]{8,18}",
            "电话号码": r"(1[3-9]\d{9}|0\d{2,3}-?\d{7,8})"
        }
        
        # 领域特定模式
        self.domain_patterns = {
            "manufacturing": {
                "压力值": r"(\d+\.?\d*)\s*(MPa|kPa|Pa)",
                "温度值": r"(\d+\.?\d*)\s*(℃|度)"
            },
            "enterprise": {
                "百分比": r"(\d+\.?\d*)\s*%",
                "日期": r"(\d{4}年\d{1,2}月\d{1,2}日|\d{4}-\d{1,2}-\d{1,2}|\d{4}/\d{1,2}/\d{1,2})",
                "部门": r"(人力|财务|技术|生产|销售|市场|研发).{0,3}(部|中心|处)"
            }
        }
        
        logger.info("增强版NER处理器初始化完成")
        logger.info(f"预定义词典: {list(self.dictionaries.keys())}")
        logger.info(f"领域词典: {list(self.domain_dictionaries.keys())}")
        logger.info(f"正则表达式模式: {list(self.patterns.keys())}")
        logger.info(f"领域模式: {list(self.domain_patterns.keys())}")
    
    def extract_entities(self, text: str, domain: str = "general") -> List[Dict[str, Any]]:
        """
        从文本中提取命名实体
        
        Args:
            text: 输入文本
            domain: 领域
            
        Returns:
            实体列表
        """
        logger.info(f"使用增强版NER处理器识别实体，输入文本: {text}，领域: {domain}")
        
        entities = []
        
        # 词典匹配
        dictionary_entities = self._extract_by_dictionary(text, domain)
        entities.extend(dictionary_entities)
        
        # 正则表达式匹配
        pattern_entities = self._extract_by_patterns(text, domain)
        entities.extend(pattern_entities)
        
        # 去重处理
        unique_entities = self._deduplicate_entities(entities)
        
        logger.info(f"识别到的实体: {unique_entities}")
        return unique_entities
    
    def _extract_by_dictionary(self, text: str, domain: str) -> List[Dict[str, Any]]:
        """
        通过词典匹配提取实体
        
        Args:
            text: 输入文本
            domain: 领域
            
        Returns:
            实体列表
        """
        entities = []
        
        # 通用词典匹配
        for entity_key, entity_values in self.dictionaries.items():
            for entity_value in entity_values:
                if entity_value in text:
                    # 查找所有匹配位置
                    start = 0
                    while True:
                        pos = text.find(entity_value, start)
                        if pos == -1:
                            break
                        entities.append({
                            "nerKey": entity_key,
                            "nerValue": entity_value,
                            "start": pos,
                            "end": pos + len(entity_value),
                            "confidence": 1.0,
                            "source": "dictionary"
                        })
                        start = pos + 1
        
        # 领域特定词典匹配
        if domain in self.domain_dictionaries:
            for entity_key, entity_values in self.domain_dictionaries[domain].items():
                for entity_value in entity_values:
                    if entity_value in text:
                        # 查找所有匹配位置
                        start = 0
                        while True:
                            pos = text.find(entity_value, start)
                            if pos == -1:
                                break
                            entities.append({
                                "nerKey": entity_key,
                                "nerValue": entity_value,
                                "start": pos,
                                "end": pos + len(entity_value),
                                "confidence": 1.0,
                                "source": "dictionary"
                            })
                            start = pos + 1
        
        logger.debug(f"通过词典匹配识别到的实体: {entities}")
        return entities
    
    def _extract_by_patterns(self, text: str, domain: str) -> List[Dict[str, Any]]:
        """
        通过正则表达式提取实体
        
        Args:
            text: 输入文本
            domain: 领域
            
        Returns:
            实体列表
        """
        entities = []
        
        # 通用模式匹配
        for entity_key, pattern in self.patterns.items():
            matches = re.finditer(pattern, text)
            for match in matches:
                entities.append({
                    "nerKey": entity_key,
                    "nerValue": match.group(0),
                    "start": match.start(),
                    "end": match.end(),
                    "confidence": 0.95,
                    "source": "pattern"
                })
        
        # 领域特定模式匹配
        if domain in self.domain_patterns:
            for entity_key, pattern in self.domain_patterns[domain].items():
                matches = re.finditer(pattern, text)
                for match in matches:
                    entities.append({
                        "nerKey": entity_key,
                        "nerValue": match.group(0),
                        "start": match.start(),
                        "end": match.end(),
                        "confidence": 0.95,
                        "source": "pattern"
                    })
        
        logger.debug(f"通过正则表达式识别到的实体: {entities}")
        return entities
    
    def _deduplicate_entities(self, entities: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """
        去除重复实体
        
        Args:
            entities: 实体列表
            
        Returns:
            去重后的实体列表
        """
        if not entities:
            return entities
        
        # 按位置排序
        entities.sort(key=lambda x: (x["start"], x["end"]))
        
        # 去重
        unique_entities = []
        unique_entities_with_pos = []  # 保留位置信息用于比较
        for entity in entities:
            # 检查是否与已添加的实体重叠
            is_duplicate = False
            for added_entity in unique_entities_with_pos:
                # 如果位置完全相同，或者一个包含另一个，则认为是重复的
                if (entity["start"] == added_entity["start"] and entity["end"] == added_entity["end"]) or \
                   (entity["start"] >= added_entity["start"] and entity["end"] <= added_entity["end"]) or \
                   (added_entity["start"] >= entity["start"] and added_entity["end"] <= entity["end"]):
                    is_duplicate = True
                    break
            
            if not is_duplicate:
                # 保留位置信息的副本用于后续比较
                unique_entities_with_pos.append(entity.copy())
                # 移除位置信息的副本用于返回结果
                entity_copy = entity.copy()
                entity_copy.pop("start", None)
                entity_copy.pop("end", None)
                unique_entities.append(entity_copy)
        
        logger.debug(f"去重后的实体: {unique_entities}")
        return unique_entities

# 测试代码
if __name__ == "__main__":
    processor = AdvancedNERProcessor()
    test_text = "承德钒钛钳工需要50MPa的压力容器操作证"
    entities = processor.extract_entities(test_text, "manufacturing")
    print("识别到的实体:")
    for entity in entities:
        print(f"  {entity['nerKey']}: {entity['nerValue']} (置信度: {entity['confidence']}, 来源: {entity['source']})")